from typing import Mapping, Dict, Hashable

import matplotlib.pyplot as plt
import nibabel
from einops import einops
from monai.config import NdarrayOrTensor
from monai.data import get_track_meta
from monai.transforms import Transform, MapTransform
from monai.config.type_definitions import DtypeLike, KeysCollection
import numpy as np
from monai.utils import TransformBackends, convert_to_tensor, convert_data_type
import cv2
import torch

"""
Use it like other transforms in MONAI,
in utils.data_utils.py, add following codes
from .expand.position_attention import position_attentionD
train_transform = transforms.Compose(
[ ...
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                Position_attentionD(keys=['image']) 
...])
"""


class positionattention(Transform):
    """
    Args:
        hidden_size: It is obviously that I am using position embedding in Transformer,
        so hidden_size is just like d_model in Transformer's paper.It should be noticed
        that hidden_size mod 6 must equals 0.
    """
    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]

    def __init__(
            self,
            dtype: DtypeLike = np.float32,
            hidden_size: int = 12,
            dims: int = 2,
            mode: str = "new"
    ) -> None:
        self.dtype = dtype
        self.hidden_size = hidden_size
        self.dims = dims
        self.mode = mode

    def build_2d_sincos_position_sym_embedding(self, grid_size, embed_dim, temperature=10000.):
        d, h, w, s = grid_size
        pad = h % 4
        if pad == 0:
            grid_h = np.concatenate(
                [np.arange(h / 4, dtype=self.dtype) / (h / 4), (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1],
                 np.arange(h / 4, dtype=self.dtype) / (h / 4), (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1]],
                axis=0)
        elif pad == 1:
            grid_h = np.concatenate(
                [np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1],
                 np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4))[::-1]],
                axis=0)
        elif pad == 2:
            grid_h = np.concatenate(
                [np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1],
                 np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1]],
                axis=0)
        elif pad == 3:
            grid_h = np.concatenate(
                [np.arange(h / 4 - 1, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1],
                 np.arange(h / 4, dtype=self.dtype) / (h / 4),
                 (np.arange(h / 4, dtype=self.dtype) / (h / 4))[::-1]],
                axis=0)
        if w % 2 == 0:
            grid_w = np.concatenate(
                [np.arange(w / 2, dtype=self.dtype) / (w / 2), (np.arange(w / 2, dtype=self.dtype) / (w / 2))[::-1]],
                axis=0)
        # 0 1 2 3 3 2 1 0
        else:
            grid_w = np.concatenate(
                [np.arange(w / 2 - 1, dtype=self.dtype) / (w / 2),
                 (np.arange(w / 2, dtype=self.dtype) / (w / 2))[::-1]],
                axis=0)
        # 0 1 2 3 4 5 6 7
        # grid_s = np.arange(s, dtype=self.dtype)
        # 0 1 2 3 4 5 6 7
        grid_w, grid_h = np.meshgrid(grid_w, grid_h)
        # print(grid_w)
        assert embed_dim % 4 == 0
        pos_dim = embed_dim // 4
        omega = np.arange(pos_dim, dtype=self.dtype) / pos_dim
        omega = 1. / (temperature ** omega)
        out_w = np.einsum('m,d->md', grid_w.flatten(), omega)
        out_h = np.einsum('m,d->md', grid_h.flatten(), omega)
        nw = out_w.reshape((h, w, pos_dim))
        nh = out_h.reshape((h, w, pos_dim))
        pos_emb = ((np.sum(np.exp(nw), axis=2) - pos_dim) + (np.sum(np.exp(nh), axis=2) - pos_dim)) * 2
        # np.save("picture.npy", pos_emb)
        pos_emb = (pos_emb + 1e-10) / pos_emb.max()
        pos_emb = np.repeat(pos_emb[:, :, None], repeats=s, axis=2)[None, :, :, :]
        pos_emb = np.repeat(pos_emb, repeats=d, axis=0)
        return pos_emb
        # out_s h np.einsum('m,d->md', [grid_s.flatten(), omega])
        # pos_emb = np.concatenate(
        #    [np.sin(out_w), np.cos(out_w), np.sin(out_h),
        #     np.cos(out_h)],
        #    axis=1)[None, :, :]
        # print(pos_emb.reshape(1, 127, 512, 12)[0, :, 0, 11])
        # pos_emb = einops.rearrange(pos_emb, "(h w s) p-> h w s p", w=w, s=s)
        # return pos_emb

    def position_attention(self, grid_size, key_point, embed_dim, temperature=10000.):
        d, h, w, s = grid_size
        center_pt_x, center_pt_y1, center_pt_y2, mid_point_y = key_point
        grid_w = np.concatenate(
            [np.arange(center_pt_x, dtype=self.dtype) / center_pt_x,
             np.arange(center_pt_x, w)[::-1] / w],
            axis=0)
        # mid_point_y = ((center_pt_y2 - center_pt_y1) // 2) + center_pt_y1
        grid_h = np.concatenate(
            [np.arange(center_pt_y1, dtype=self.dtype) / center_pt_y1,
             np.arange(center_pt_y1, mid_point_y, dtype=self.dtype)[::-1] / mid_point_y,
             np.arange(mid_point_y, center_pt_y2, dtype=self.dtype) / center_pt_y2,
             np.arange(center_pt_y2, h, dtype=self.dtype)[::-1] / h],
            axis=0)

        grid_w, grid_h = np.meshgrid(grid_w, grid_h)
        assert embed_dim % 4 == 0
        pos_dim = embed_dim // 4
        omega = np.arange(pos_dim, dtype=self.dtype) / pos_dim
        omega = 1. / (temperature ** omega)
        out_w = np.einsum('m,d->md', grid_w.flatten(), omega)
        out_h = np.einsum('m,d->md', grid_h.flatten(), omega)
        nw = out_w.reshape((h, w, pos_dim))
        nh = out_h.reshape((h, w, pos_dim))
        pos_emb = ((np.sum(np.exp(nw), axis=2) - pos_dim) + (np.sum(np.exp(nh), axis=2) - pos_dim)) * 2
        # np.save("picture.npy", pos_emb)
        # print(pos_emb)
        pos_emb = pos_emb / pos_emb.max()
        pos_emb = np.repeat(pos_emb[:, :, None], repeats=s, axis=2)[None, :, :, :]
        pos_emb = np.repeat(pos_emb, repeats=d, axis=0)
        return pos_emb

    def position_attention_v2(self, grid_size, key_point, embed_dim, temperature=10000.):
        d, h, w, s = grid_size
        min_x, max_x, min_y, max_y = key_point
        y = max_y - min_y
        pad = y % 4
        if pad == 0:
            grid_h = np.concatenate(
                [np.zeros(min_y, dtype=self.dtype),
                 np.arange(y / 4, dtype=self.dtype) / (y / 4), (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1],
                 np.arange(y / 4, dtype=self.dtype) / (y / 4), (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1],
                 np.zeros(h - max_y, dtype=self.dtype)],
                axis=0)

        elif pad == 1:
            grid_h = np.concatenate(
                [np.zeros(min_y, dtype=self.dtype), np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1],
                 np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4))[::-1], np.zeros(h - max_y, dtype=self.dtype)],
                axis=0)

        elif pad == 2:
            grid_h = np.concatenate(
                [np.zeros(min_y, dtype=self.dtype), np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1],
                 np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1], np.zeros(h - max_y, dtype=self.dtype)],
                axis=0)

        elif pad == 3:
            grid_h = np.concatenate(
                [np.zeros(min_y, dtype=self.dtype), np.arange(y / 4 - 1, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1],
                 np.arange(y / 4, dtype=self.dtype) / (y / 4),
                 (np.arange(y / 4, dtype=self.dtype) / (y / 4))[::-1], np.zeros(h - max_y, dtype=self.dtype)],
                axis=0)

        x = max_x-min_x
        if x % 2 == 0:
            grid_w = np.concatenate(
                [np.zeros(min_x),
                 np.arange(x / 2, dtype=self.dtype) / (x / 2),
                 (np.arange(x / 2, dtype=self.dtype) / (x / 2))[::-1],
                 np.zeros(w-max_x)],
                axis=0)
        # 0 1 2 3 3 2 1 0
        else:
            grid_w = np.concatenate(
                [np.zeros(min_x),
                 np.arange(x / 2 - 1, dtype=self.dtype) / (x / 2),
                 (np.arange(x / 2, dtype=self.dtype) / (x / 2))[::-1],
                 np.zeros(w - max_x)],
                axis=0)
        #print(np.arange(center_pt_x, dtype=self.dtype).shape, np.arange(center_pt_x)[::-1][:w - center_pt_x].shape)
        # mid_point_y = ((center_pt_y2 - center_pt_y1) // 2) + center_pt_y1
        """
        grid_h = np.concatenate(
            [np.arange(center_pt_y1, dtype=self.dtype) / center_pt_y1,
             np.arange(center_pt_y1, dtype=self.dtype)[::-1][:mid_point_y - center_pt_y1] / center_pt_y1,
             np.arange(center_pt_y2, dtype=self.dtype)[mid_point_y:] / center_pt_y2,
             np.arange(center_pt_y2, dtype=self.dtype)[::-1][:h - center_pt_y2] / center_pt_y2],
            axis=0)
        """

        grid_w, grid_h = np.meshgrid(grid_w, grid_h)
        assert embed_dim % 4 == 0
        pos_dim = embed_dim // 4
        omega = np.arange(pos_dim, dtype=self.dtype) / pos_dim
        omega = 1. / (temperature ** omega)
        out_w = np.einsum('m,d->md', grid_w.flatten(), omega)
        out_h = np.einsum('m,d->md', grid_h.flatten(), omega)
        nw = out_w.reshape((h, w, pos_dim))
        nh = out_h.reshape((h, w, pos_dim))
        pos_emb = ((np.sum(np.exp(nw), axis=2) - pos_dim) + (np.sum(np.exp(nh), axis=2) - pos_dim)) * 2
        # np.save("picture.npy", pos_emb)
        # print(pos_emb)
        pos_emb = pos_emb / pos_emb.max()
        pos_emb = np.repeat(pos_emb[:, :, None], repeats=s, axis=2)[None, :, :, :]
        pos_emb = np.repeat(pos_emb, repeats=d, axis=0)
        return pos_emb

    def find_coord(self, img, kernel_size=13):
        hard_mask = img <= img.mean()
        hard_mask = hard_mask.astype("uint8")
        # print(hard_mask.shape)
        dilate_kernel = np.ones((kernel_size, kernel_size), dtype="uint8")
        hard_mask = cv2.erode(hard_mask, dilate_kernel)
        coords = np.array(np.where(hard_mask == 0))
        coords_x = np.mean(coords[0, :])
        coords_y = np.mean(coords[1, :])
        # coords_mid_y = np.median(coords[1, :])
        coords_y_length = np.std(coords[1, :])
        center_pt_x = round(coords_x)
        # center_pt_x = round(center_pt_x)
        center_pt_y1, center_pt_y2 = round(coords_y - coords_y_length / 2), round(coords_y + coords_y_length / 2)
        key_point = (center_pt_x, center_pt_y1, center_pt_y2)
        return key_point

    def find_coord_v2(self, img, kernel_size=20):
        hard_mask = img > img.mean()
        hard_mask = hard_mask.astype("uint8")
        # np.save('origin_mid.npy', hard_mask)
        dilate_kernel = np.ones((kernel_size, kernel_size), dtype="uint8")
        hard_mask = cv2.erode(hard_mask, dilate_kernel)
        # np.save("mid.npy", hard_mask)
        coords = np.where(hard_mask == 1)
        coords_y1, coords_y2 = np.min(coords[0]), np.max(coords[0])  # 行标
        coords_x1, coords_x2 = np.min(coords[1]), np.max(coords[1])  # 列标 boundary
        # coords_mid_y = (coords_y2 - coords_y1) // 2 + coords_y1
        # center_pt_x = (coords_x2 - coords_x1) // 2 + coords_x1
        # np.save('coordx.npy', np.array(coord
        # center_pt_x = round(center_pt_x)
        # center_pt_y1, center_pt_y2 = round((coords_mid_y - coords_y1) / 2), round(
        # (coords_y2 - coords_mid_y) / 2 + coords_mid_y)

        key_point = (coords_x1, coords_x2, coords_y1, coords_y2)
        # print(key_point)
        return key_point

    def __call__(self, img: NdarrayOrTensor, *args, **kwargs) -> NdarrayOrTensor:
        """
                Apply the transform to `img`.
        """
        img = convert_to_tensor(img, track_meta=get_track_meta())
        img_size = img.shape

        dtype = self.dtype or img.dtype
        if self.mode == "old":
            pos_emb = self.build_2d_sincos_position_sym_embedding(
                img_size,
                self.hidden_size)
        elif self.mode == "new":
            key_point = self.find_coord_v2(np.squeeze(img, axis=0))  # Note that, if B > 1, then error
            pos_emb = self.position_attention_v2(
                img_size,
                key_point,
                self.hidden_size
            )
        # np.save("original.npy", img)
        img = img * pos_emb
        # np.save("processing.npy", img)
        ret: NdarrayOrTensor = convert_data_type(img, dtype=dtype)[0]

        return ret


class PositionattentionD(MapTransform):
    backend = positionattention.backend

    def __init__(
            self,
            keys: KeysCollection,
            hidden_size: int = 12,
            dtype: DtypeLike = np.float32,
            allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.syscoding = positionattention(hidden_size=hidden_size, dtype=dtype)

    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.syscoding(d[key])
        return d

if __name__ == '__main__':
    test = PositionattentionD(keys='img')
    img = nibabel.load("/home/usr/xulongwei/dataset/dataset1/imagesTr/coronacases_org_001.nii")
    img = img.get_fdata()
    img = img[None, :, :, :]
    label = np.zeros_like(img)
    test_a = positionattention()
    test_d = {'img': img, 'label': label}
    out = test_a(img)
    # out = test(test_d)
    # out = out.type_as(img_size)
    # print(torch.equal(out['img'], img))
    # print(torch.equal(out['label'], label))
